from __future__ import annotations
from typing import Optional, TYPE_CHECKING
from .base_structure import Transaction
from itertools import combinations
from math import ceil

if TYPE_CHECKING:
    from .base_structure import Item


def get_item_frequency(data: list[Transaction]) -> dict[Item, int]:
    item_frequency: dict[Item, int] = {}
    for transaction in data:
        for item in transaction:
            if item in item_frequency:
                item_frequency[item] += transaction.weight
            else:
                item_frequency[item] = transaction.weight
    return item_frequency


class Abstract_Node:
    def __init__(self) -> None:
        self.children: dict[Item, Node] = {}

    def add_item(self, item: Item, value: int = 1) -> Node:
        child = Node(item).set_parent(self).add_value(value)
        self.children[item] = child
        return child

    def add_child(self, child: Node) -> Node:
        self.children[child.item] = child
        return child

    def find_item(self, item: Item) -> Optional[Node]:
        if item in self.children:
            return self.children[item]
        else:
            return None


class Root(Abstract_Node):
    def __init__(self) -> None:
        super().__init__()


class Node(Abstract_Node):
    def __init__(self, item: Item) -> None:
        self.item = item
        self.value = 0
        super().__init__()

    def set_parent(self, parent: Abstract_Node) -> Node:
        parent.add_child(self)
        self.parent = parent
        return self

    def add_value(self, n: int) -> Node:
        self.value += n
        return self


class FPtree:
    def __init__(self, data: list[Transaction], item_frequency: dict[Item, int], threshold: int) -> None:
        self.threshold = threshold
        self.root = Root()
        self.item_nodes_map: dict[Item, list[Node]] = {}

# form tree
        for transaction in data:
            transaction = transaction.filter(item_frequency, threshold)
            transaction = transaction.ordered(item_frequency)
            weight = transaction.weight
            node: Abstract_Node = self.root
            for item in transaction:
                _node = node.find_item(item)
                if _node is None:
                    node = node.add_item(item, weight)
                    self.item_nodes_map.setdefault(item, []).append(node)
                else:
                    node = _node.add_value(weight)

        self.item_frequency = {item: sum(
            [node.value for node in nodes]) for item, nodes in self.item_nodes_map.items()}

# 顺序排列item
        self.ordered_items = sorted(self.item_frequency.keys(
        ), key=lambda x: (self.item_frequency[x], x.item_name))

    def form_conditional_tree(self, item: Item) -> FPtree:
        data: list[Transaction] = []
        for node in self.item_nodes_map[item]:
            weight = node.value
            transaction = Transaction(weight)
            parent = node.parent
            while isinstance(parent, Node):
                transaction.append(parent.item)
                parent = parent.parent
            data.append(transaction)
        item_frequency = get_item_frequency(data)
        return FPtree(data, item_frequency, self.threshold)


    def is_empty(self) -> bool:
        return len(self.root.children) == 0

    def is_single_path(self) -> bool:
        for _, nodes in self.item_nodes_map.items():
            if len(nodes) > 1:
                return False
        return True


class FPminer:
    def __init__(self, data: list[Transaction], support: float = 1.0, threshold: Optional[int] = None) -> None:
        self.data = data
        self.len = len(data)
        if threshold is None:
            if support < 0 or support > 1:
                raise ValueError(
                    f'Support must be float between 0 and 1, got {support}')
            else:
                self.threshold = ceil(self.len * support)
        else:
            if threshold < 0 or threshold > self.len:
                raise ValueError(
                    f'Threshold must be value between 0 and {self.len}, got {threshold}')
            self.threshold = int(threshold)

    def get_item_frequency(self) -> dict[Item, int]:
        return get_item_frequency(self.data)

    def get_frequent_patterns(self) -> dict[frozenset[Item], int]:
        item_frequency = self.get_item_frequency()
        tree = FPtree(self.data, item_frequency, self.threshold)
        return self.conditional_tree_miner(set(), tree, {})

    # 递归获取频繁模式
    def conditional_tree_miner(self, conditional_pattern_base: set[Item], conditional_tree: FPtree, frequent_patterns: dict[frozenset[Item], int]) -> dict[frozenset[Item], int]:
        frequent_patterns = frequent_patterns
        # 空树直接返回
        if conditional_tree.is_empty():
            return frequent_patterns
        # 无分支树，直接获取幂集与条件模式基的并集
        elif conditional_tree.is_single_path():
            frequency = conditional_tree.item_frequency
            for i in range(1, len(conditional_tree.ordered_items) + 1):
                for conditional_pattern in combinations(conditional_tree.ordered_items, i):
                    value = min([frequency[item] for item in conditional_pattern])
                    frequent_patterns[frozenset(conditional_pattern_base | set(conditional_pattern))] = value
            return frequent_patterns
        # 其余遍历条件树中的项，与原条件模式基取并集为新条件模式基并递归
        else:
            for item in conditional_tree.ordered_items:
                new_conditional_pattern_base = conditional_pattern_base | {item}
                new_conditional_tree = conditional_tree.form_conditional_tree(item)
                frequent_patterns[frozenset(new_conditional_pattern_base)] = conditional_tree.item_frequency[item]
                frequent_patterns |= self.conditional_tree_miner(new_conditional_pattern_base, new_conditional_tree, frequent_patterns)
            return frequent_patterns
